# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import abc
import argparse
import datetime as dt
import logging
from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence, cast

from typing_extensions import override

from crossbench.benchmarks.base import SubStoryBenchmark
from crossbench.benchmarks.embedder.config.cujs import CUJsConfig
from crossbench.benchmarks.embedder.config.setup_commands import \
    SetupCommandsConfig
from crossbench.cli.ui import timer
from crossbench.parse import ObjectParser
from crossbench.stories.story import Story

if TYPE_CHECKING:
  from crossbench.action_runner.config import ActionRunnerConfig
  from crossbench.benchmarks.loading.config.blocks import ActionBlock
  from crossbench.browsers.webview.embedder import WebviewEmbedder
  from crossbench.cli.parser import CrossBenchArgumentParser
  from crossbench.cli.types import Subparsers
  from crossbench.runner.run import Run


class EmbedderStory(Story, metaclass=abc.ABCMeta):

  def __init__(self, name: str, actions: Sequence[ActionBlock]):
    duration = dt.timedelta(seconds=30)
    self._actions = actions
    super().__init__(name, duration)

  def setup(self, run: Run) -> None:
    # TODO(zbikowski): Add a way to ensure embedder is installed.
    # Launching the Google Quick Search app
    run_benchmark = cast(EmbedderBenchmark, run.benchmark)
    if run_benchmark.embedder_drop_caches:
      run.browser_platform.sh("am", "kill-all")
      run.browser_platform.sh(  # noqa: S604
          "echo 3 > /proc/sys/vm/drop_caches", shell=True)
    run_browser = cast("WebviewEmbedder", run.browser)
    run.browser_platform.sh("am", "start", "-S", "-W", "-a",
                            f"{run_browser.android_package}.GOOGLE_SEARCH",
                            "-n",
                            f"{run_browser.android_package}/.SearchActivity")
    logging.info("Starting Embedder Benchmark...")

  def run(self, run: Run) -> None:
    with timer():
      logging.info("-" * 80)
      action_runner = run.action_runner
      for block in self._actions:
        for action in block:
          action.run_with(run, action_runner)
    # Empty line to preserve timer output.
    print()
    logging.info("Stopping Embedder Benchmark...")

  def teardown(self, run: Run) -> None:
    # Return to home screen after the story has run
    run.browser_platform.sh("am", "start", "-a", "android.intent.action.MAIN",
                            "-c", "android.intent.category.HOME")

  @classmethod
  def all_story_names(cls) -> tuple[str, ...]:
    return ()


class EmbedderBenchmark(SubStoryBenchmark):
  """
  Benchmark runner for a WV embedder app mode.
  """
  NAME: ClassVar = "embedder"
  DEFAULT_STORY_CLS: ClassVar = EmbedderStory

  def __init__(
      self,
      stories: Sequence[Story],
      action_runner_config: Optional[ActionRunnerConfig] = None,
      embedder_process_name: str = "search",
      embedder_setup_command_config: Optional[SetupCommandsConfig] = None,
      embedder_drop_caches: bool = False
  ) -> None:
    for story in stories:
      assert isinstance(story, self.DEFAULT_STORY_CLS)
    super().__init__(stories, action_runner_config)
    self._embedder_process_name = embedder_process_name
    self._embedder_setup_command_config = embedder_setup_command_config
    self._embedder_drop_caches = embedder_drop_caches

  @property
  def embedder_process_name(self) -> str:
    return self._embedder_process_name

  @property
  def embedder_setup_command_config(self) -> Optional[SetupCommandsConfig]:
    return self._embedder_setup_command_config

  @property
  def embedder_drop_caches(self) -> bool:
    return self._embedder_drop_caches

  @classmethod
  @override
  def cli_description(cls) -> str:
    assert cls.__doc__
    return cls.__doc__.strip()

  @classmethod
  @override
  def add_cli_parser(cls, subparsers: Subparsers) -> CrossBenchArgumentParser:
    parser = super().add_cli_parser(subparsers)
    parser.add_argument(
        "--cujs-config",
        "--cuj-config",
        type=CUJsConfig.parse,
        help="Stories we want to perform in the benchmark run following a"
        "specified scenario.")
    parser.add_argument(
        "--embedder-process-name",
        default="search",
        help="Name of the embedder process.")
    parser.add_argument(
        "--embedder-setup-command-config",
        type=SetupCommandsConfig.parse,
        help="Setup commands to run before the benchmark.")
    parser.add_argument(
        "--embedder-drop-caches",
        default=False,
        action="store_true",
        help="Drop page cache before running each story.")
    return parser

  @classmethod
  @override
  def stories_from_cli_args(cls, args: argparse.Namespace) -> Sequence[Story]:
    config = cls.get_cujs_config(args)
    cujs = tuple(
        EmbedderStory(
            name=cuj_config.label,
            actions=cuj_config.blocks,
        ) for cuj_config in config.cujs)
    return cujs

  @classmethod
  @override
  def kwargs_from_cli(cls, args: argparse.Namespace) -> dict[str, Any]:
    kwargs = super().kwargs_from_cli(args)
    kwargs["embedder_process_name"] = args.embedder_process_name
    kwargs["embedder_setup_command_config"] = args.embedder_setup_command_config
    kwargs["embedder_drop_caches"] = args.embedder_drop_caches
    return kwargs

  @classmethod
  def get_cujs_config(cls, args: argparse.Namespace) -> CUJsConfig:
    if global_config := args.config:
      # TODO: migrate --config to an already parsed hjson/json dict
      config_file = global_config
      config_data = ObjectParser.hjson_file(config_file)
      if cujs_config_dict := config_data.get("cujs"):
        if args.cujs_config:
          raise argparse.ArgumentTypeError(
              "Conflicting arguments: "
              "either specify a --config file without a 'cujs' property "
              "or remove the --cuj-config argument.")
        # TODO: CUJsConfig.parse_dict should be able to parse the inner dict.
        return CUJsConfig.parse_dict({"cujs": cujs_config_dict})
    return args.cujs_config
